# std
import os
import datetime
import functools
from collections import namedtuple
import traceback

# 3rd party
import pandas as pd
from absl import app, flags

# internal
from coin.proto.coin_order_gateway_pb2 import OrderEvent
import coin.proto.coin_order_enums_pb2 as order_pb2
from coin.proto.support.coin_order_latency_pb2 import OrderLatency
from coin.proto.support.coin_order_stateless_pb2 import StatelessOrderEvent

import coin.strategy.mm.tool.archive_base as abase
from coin.exchange.kr_rest.product.product_impl import generate_product_from_str
from coin.flow.topic_record import TopicRecordWriter

FEED_MACHINES = [
    'feed-best.ap-northeast-1.aws',
    'feed-best.cn-hongkong.aliyun',
    # 'feed-01.ap-northeast-2.aws',
    # 'feed-01.eu-west-1.aws',
    # 'feed-01.us-east-1.aws',
    # 'feed-02.us-east-1.aws'
]

ORDER_ROOT_DIR = '/remote/iosg/coin/data/og_proto_stateless_log'
FEED_ROOT_DIR = None
FEED_CACHE_DIR = '/remote/iosg/coin-2/feed_cache'

import getpass

if getpass.getuser() in ["jhkim"]:
  OUT_DIR = f'/remote/iosg/home-2/{getpass.getuser()}/order_latency'
else:
  OUT_DIR = './order_latency'

flags.DEFINE_string('date', None, 'example: --date 20190708')
flags.mark_flag_as_required('date')
flags.register_validator('date', lambda value: value.isdigit(), message='date example: 20190708')

Bbo = namedtuple('Bbo', ('timestamp', 'price', 'qty'))

# bbo history for one symbol


class Product(object):
  def __init__(self):
    super().__init__()
    self.bbo_asks = []
    self.bbo_bids = []

  def update(self, book_event):
    book = book_event
    ask = book.ask0
    bid = book.bid0

    if ask:
      bbo_ask = Bbo(timestamp=book.timestamp, price=ask.price, qty=ask.qty)
      if not self.bbo_asks or self.bbo_asks[-1][1:] != bbo_ask[1:]:
        self.bbo_asks.append(bbo_ask)

    if bid:
      bbo_bid = Bbo(timestamp=book.timestamp, price=bid.price, qty=bid.qty)
      if not self.bbo_bids or self.bbo_bids[-1][1:] != bbo_bid[1:]:
        self.bbo_bids.append(bbo_bid)


def get_order_submits(market_type, exchange, date, machine, owner):

  start_time = datetime.datetime.strptime(date, '%Y%m%d').replace(tzinfo=datetime.timezone.utc)
  end_time = start_time + datetime.timedelta(days=1)

  start_ts = int(
      datetime.datetime.timestamp(start_time.replace(tzinfo=datetime.timezone.utc)) * 1e9)
  end_ts = int(datetime.datetime.timestamp(end_time.replace(tzinfo=datetime.timezone.utc)) * 1e9)

  orders = []

  def on_log(timestamp, raw):
    if timestamp < start_ts or timestamp >= end_ts:
      return

    log = StatelessOrderEvent()
    try:
      log.ParseFromString(raw)
      if log.event.type == OrderEvent.ORDER_ACCEPTED:
        orders.append(log)
    except Exception as e:
      print('og log parsing error: ', e)
      return

  abase.run_from_stateless_order_archive(on_log=on_log,
                                         market_type=market_type,
                                         exchange=exchange,
                                         machine=machine,
                                         root_dir=ORDER_ROOT_DIR,
                                         owner=owner,
                                         start_time=start_time,
                                         end_time=end_time)

  return orders  # list of StatelessOrderEvent


BookEvent = namedtuple('BookEvent', ['product', 'timestamp', 'bid0', 'ask0'])


def get_feed_stream(date, products, machine) -> pd.DataFrame:
  start_time = datetime.datetime.strptime(date, '%Y%m%d')
  end_time = start_time + \
      datetime.timedelta(days=1) + \
      datetime.timedelta(minutes=35 + 35)  # +-35 mins padding

  feeds = []

  def on_book(product, book):
    feeds.append(BookEvent(product, book.timestamp, book.bid0(), book.ask0()))

  def on_trade(product, trade):
    # feeds.append((product, trade.timestamp, 'trade', trade))
    pass

  def on_book_reset(_book_builder_name, book_builder):
    for product in products:
      book_builder.subscribe(product, functools.partial(on_book, product))
      if hasattr(book_builder, "subscribe_trade"):
        book_builder.subscribe_trade(product, functools.partial(on_trade, product))

  abase.run_from_feed_cache_base(
      products,
      on_book_reset,
      start_time,
      end_time,
      machine,
      worker_ids=['1', '2'],
      write_feed_cache_if_none=False,
      feed_cache_dir=FEED_CACHE_DIR,
  )
  return feeds


# HACK, copied from archive_base.py
# this script does not use baseflags
def order_symbols_to_products(symbols, date, market_type, exchange):
  # .replace(tzinfo=datetime.timezone.utc)
  start_time = datetime.datetime.strptime(date, '%Y%m%d')

  products = []
  for symbol in symbols:
    if market_type == 'Futures':
      product = generate_product_from_str(exchange, 'futures', symbol, start_time)
    else:
      product = generate_product_from_str(exchange, None, symbol)
    products.append(product)

  return products


def region_from_machine(m):
  return m.split('.', 3)[1]


def main(_):
  date = flags.FLAGS.date
  cols = ['machine', 'market_type', 'exchange', 'owner']
  print('reading sql ..')
  df = pd.read_sql(
      f'select {",".join(cols)} from StrategyPtaStatsView '
      f'where type="SYMBOL" and trading_date="{date}";',
      'mysql://pta_querier:0C3n5szBUEhn@coin-db.corp.prestolabs.io/pta_stats_20190513',
  )
  df.drop_duplicates(inplace=True)
  df.reset_index(drop=True, inplace=True)
  Record = namedtuple('Record',
                      [
                          'order_machine',
                          'feed_machine',
                          'market_type',
                          'exchange',
                          'symbol',
                          'timestamp',
                          'price',
                          'qty',
                          'latency'
                      ])
  records = []

  for (index, row) in df.iterrows():
    print('------')
    print(row)

    try:
      strategy_region = region_from_machine(row.machine)
      if "us-east-2" in strategy_region:
        strategy_region = strategy_region.replace("us-east-2", "us-east-1")

      if strategy_region == 'aliyun-cn-hongkong':
        strategy_region = 'cn-hongkong'
      feed_machines = list(
          filter(lambda m: region_from_machine(m) == strategy_region, FEED_MACHINES))

      if not feed_machines:
        # e.g. Kraken
        print('no feed machine for strategy-region', strategy_region)
        continue

      # list of StatelessOrderEvent (Accept)
      orders = get_order_submits(row.market_type, row.exchange, date, row.machine, row.owner)
      print(f'{len(orders)} orders')

      symbols = set()
      for order in orders:
        symbols.add(order.event.symbol)
      symbols = list(symbols)

      products = order_symbols_to_products(symbols, date, row.market_type, row.exchange)
      order_symbol_to_product_map = {
          symbol: product for (symbol, product) in zip(symbols, products)
      }
      assert len(feed_machines) == 1, f'more than feed machines: {feed_machines}'
      feed_machine = feed_machines[0]
      # book + trade => book only
      feeds = get_feed_stream(date, products, feed_machine)
      print(f'{len(feeds)} feeds')

      i_order = 0
      i_feed = 0

      # trace = {}
      # order_by_oid = {}
      bbos = {product: Product() for product in products}

      # (X) order.submit_time < book.time < order.accept_time
      # (O) order.submit_time < book.time < order.accept_time + 10 secs (10*1e9 ns)

      outs = []  # [(submit_timestamp, OrderLatency)]

      # iterate two streams by time
      while i_order < len(orders) or i_feed < len(feeds):
        # order
        if i_feed == len(feeds) or (
            i_order < len(orders)
            and orders[i_order].event.event_time + 10 * 1e9 < feeds[i_feed].timestamp):

          order = orders[i_order]
          i_order += 1

          if order.state.order_submitted_time > 0 and order.state.order_accepted_time > 0:
            product = order_symbol_to_product_map[order.event.symbol]
            if order.event.order_side in [
                order_pb2.BUY_ORDER, order_pb2.BUY_OPEN_ORDER, order_pb2.BUY_CLOSE_ORDER
            ]:
              # look for bid items
              feed_history = bbos[product].bbo_bids
            elif order.event.order_side in [
                order_pb2.SELL_ORDER, order_pb2.SELL_OPEN_ORDER, order_pb2.SELL_CLOSE_ORDER
            ]:
              feed_history = bbos[product].bbo_asks
            else:
              # unknown event
              raise RuntimeError('[error] unknown order side')

            latency = None

            # scan backward
            for i in range(len(feed_history) - 1, 0, -1):  # i>=1
              if feed_history[i].timestamp < order.state.order_submitted_time:
                break
              if (feed_history[i].price == order.event.order_price
                  and feed_history[i].qty - feed_history[i - 1].qty == order.event.order_qty
                  and feed_history[i].price == feed_history[i - 1].price):
                # if price changed, we don't know diff qty
                latency = feed_history[i].timestamp - \
                    order.state.order_submitted_time
                break

            if latency:
              out = OrderLatency()
              out.internal_order_id = order.event.internal_order_id
              if order.event.external_order_id:
                out.external_order_id = order.event.external_order_id
              out.latency_ns = latency
              outs.append((order.state.order_submitted_time, out.SerializeToString()))

              records.append(
                  Record(
                      order_machine=row.machine,
                      feed_machine=feed_machine,
                      market_type=row.market_type,
                      exchange=row.exchange,
                      symbol=product.symbol,
                      timestamp=order.state.order_submitted_time,
                      price=order.event.order_price,
                      qty=order.event.order_qty,
                      latency=latency,
                  ))

        # feed
        else:
          feed = feeds[i_feed]
          i_feed += 1
          bbos[feed.product].update(feed)

      path = f'{OUT_DIR}/{row.machine}/{date}/{row.owner}.{row.market_type}.{row.exchange}.gz'
      os.makedirs(os.path.dirname(path), exist_ok=True)

      if os.path.exists(path):
        os.remove(path)
      with TopicRecordWriter(path) as writer:
        for out in outs:
          writer.write(None, 0, None, out[0], out[1])

    except Exception as e:
      print(e)
      traceback.print_tb(e.__traceback__)

  print(f'converting {len(records)} records to DataFrame ..')
  df = pd.DataFrame(records)
  print('saving ..')
  os.makedirs(f'{OUT_DIR}/dataframes/', exist_ok=True)
  df.to_hdf(f'{OUT_DIR}/dataframes/{date}.h5', 'df')


if __name__ == "__main__":
  app.run(main)
